package com.nstc.sftpExchange.config.sftp;

import cn.hutool.core.collection.CollectionUtil;
import com.google.common.cache.CacheBuilder;
import com.google.common.cache.CacheLoader;
import com.google.common.cache.LoadingCache;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.pool2.impl.GenericObjectPool;
import org.apache.commons.pool2.impl.GenericObjectPoolConfig;
import org.springframework.beans.BeanUtils;
import org.springframework.beans.factory.InitializingBean;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ExecutionException;

@Slf4j
public class SftpRoutingPool implements InitializingBean {

    private final SftpConfigProperties sftpPoolProperties;

    private final Map<FtpConnectionKey, SftpConfigProperties.SftpConfig> poolPropertiesMap;

    private LoadingCache<FtpConnectionKey, GenericObjectPool<FtpConnection>> poolCache;

    public SftpRoutingPool(SftpConfigProperties sftpPoolProperties) {
        this.sftpPoolProperties = sftpPoolProperties;
        this.poolPropertiesMap = new HashMap<>();
    }

    public GenericObjectPool<FtpConnection> getPool(FtpConnectionKey connectionKey) {
        GenericObjectPool<FtpConnection> pool = null;
        try {
            pool = this.poolCache.get(connectionKey);
        } catch (ExecutionException e) {
            throw new RuntimeException("sftp获取连接失败", e);
        }
        return pool;
    }

    @Override
    public void afterPropertiesSet() throws Exception {
        List<SftpConfigProperties.SftpConfig> sftpConfigs = sftpPoolProperties.getSftp();
        if (CollectionUtil.isNotEmpty(sftpConfigs)) {
            for (SftpConfigProperties.SftpConfig sftpConfig : sftpConfigs) {
                FtpConnectionKey connectionKey = new FtpConnectionKey();
                connectionKey.setHost(sftpConfig.getHost());
                connectionKey.setPort(sftpConfig.getPort());
                connectionKey.setUsername(sftpConfig.getUsername());
                connectionKey.setName(sftpConfig.getName());
                poolPropertiesMap.put(connectionKey, sftpConfig);
            }
        }
        poolCache = CacheBuilder.newBuilder()
                .build(new CacheLoader<FtpConnectionKey, GenericObjectPool<FtpConnection>>() {
                    @Override
                    public GenericObjectPool<FtpConnection> load(FtpConnectionKey connectionKey) throws Exception {
                        GenericObjectPool<FtpConnection> pool = null;
                        SftpProperties sftpInfo = new SftpProperties();

                        sftpInfo.setProtocol(connectionKey.getProtocol());
                        sftpInfo.setHost(connectionKey.getHost());
                        sftpInfo.setPort(connectionKey.getPort());
                        sftpInfo.setUsername(connectionKey.getUsername());
                        sftpInfo.setPassword(connectionKey.getPassword());
                        sftpInfo.setPrivateKeyPath(connectionKey.getPrivateKeyPath());

                        SftpConfigProperties.SftpConfig poolProperties = poolPropertiesMap.get(connectionKey);
                        if (poolPropertiesMap.containsKey(connectionKey)) {
                            sftpInfo.setConnectTimeout(poolProperties.getConnectTimeout());
                            sftpInfo.setReadTimeout(poolProperties.getReadTimeout());

                            GenericObjectPoolConfig poolConfig = new GenericObjectPoolConfig();
                            BeanUtils.copyProperties(poolConfig, poolProperties.getPool());
                            poolConfig.setMinIdle(poolProperties.getPool().getMinIdle());
                            poolConfig.setMaxIdle(poolProperties.getPool().getMaxIdle());
                            poolConfig.setMaxTotal(poolProperties.getPool().getMaxTotal());
                            poolConfig.setMaxWaitMillis(poolProperties.getPool().getMaxWaitMillis());
                            poolConfig.setJmxEnabled(false);
                            sftpInfo.setPool(poolConfig);
                        } else {
                            sftpInfo.setConnectTimeout(poolProperties.getConnectTimeout());
                            sftpInfo.setReadTimeout(poolProperties.getReadTimeout());
                            sftpInfo.setPool(poolProperties.getPool());
                        }
                        return new GenericObjectPool<>(new SftpConnectionFactory(sftpInfo), sftpInfo.getPool());
                    }
                });
    }
}
